import * as tf from '@tensorflow/tfjs'
import * as tfvis from '@tensorflow/tfjs-vis'

window.onload = async () => {
  const heights = [150, 160, 170]
  const weights = [40, 50, 60]

  tfvis.render.scatterplot(
    { name: '身高体重训练数据' },
    { values: heights.map((x, i) => ({ x, y: weights[i] })), series: ['身高体重'] },
    {
      xAxisDomain: [140, 180],
      yAxisDomain: [30, 70],
    }
  )

  // 归一化数据为 [0,1] 之间
  const inputs = tf.tensor(heights).sub(150).div(20)
  const labels = tf.tensor(weights).sub(40).div(20)

  // 训练预测
  // 初始化模型
  const model = tf.sequential()
  // 添加层
  model.add(tf.layers.dense({ units: 1, inputShape: [1] }))
  // 损失函数，优化器
  model.compile({
    loss: tf.losses.meanSquaredError, // 损失函数
    optimizer: tf.train.sgd(0.1), // 优化器,学习率
  })

  await model.fit(inputs, labels, {
    batchSize: 3,
    epochs: 200,
    callbacks: tfvis.show.fitCallbacks(
      {
        name: '训练过程',
      },
      ['loss']
    ),
  })

  let guessX = 180
  const output = model.predict(tf.tensor([guessX]).sub(150).div(20))

  output.print()
  console.log(`如果身高为 ${guessX}, 输出体重为 ${output.mul(20).add(40).dataSync()[0]}`)
}
